from pymilvus import model, MilvusClient


def test1():
    docs = [
        "Artificial intelligence was founded as an academic discipline in 1956.",
        "Alan Turing was the first person to conduct substantial research in AI.",
        "Born in Maida Vale, London, Turing was raised in southern England.",
    ]

    sentence_transformer_ef = model.dense.SentenceTransformerEmbeddingFunction(
        # model_name='all-MiniLM-L6-v2',
        device='cpu'
    )

    vectors = sentence_transformer_ef.encode_documents(docs)
    data = [{"id": i, "vector": vectors[i], "text": docs[i]} for i in range(len(vectors))]



def test2():
    embedding_fn = model.DefaultEmbeddingFunction()

    docs = [

        "人工智能作为一门学科在1956年成立。",

        "艾伦·图灵是最早进行实质性人工智能研究的人之一。",

        "图灵出生在伦敦的梅达维尔，并在英格兰南部长大。",

    ]

    vectors = embedding_fn.encode_documents(docs)

    # 输出的向量有768个维度，正好与我们刚刚创建的集合相匹配。

    print("维度:", embedding_fn.dim, vectors[0].shape)  # 维度: 768 (768,)

    # 每个实体都有id，向量表示，原始文本和主题标签。

    data = [

        {"id": i, "vector": vectors[i], "text": docs[i], "subject": "历史"}

        for i in range(len(vectors))

    ]

    print("数据包含", len(data), "个实体，每个实体包含的字段为：", data[0].keys())

    print("向量维度：", len(data[0]["vector"]))

    # 创建一个集合并插入数据

    client = MilvusClient('./milvus_local.db')

    client.create_collection(

        collection_name="demo_collection",

        dimension=768,  # 用于此次演示的向量具有768个维度

    )

    client.insert(collection_name="demo_collection", data=data)

if __name__ == '__main__':
    test1()